// Copyright 2020 The XLS Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Bindings class (name to defining AST node tracking) for use in parsing.

#ifndef XLS_DSLX_FRONTEND_BINDINGS_H_
#define XLS_DSLX_FRONTEND_BINDINGS_H_

#include <algorithm>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <variant>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "xls/dslx/frontend/ast.h"
#include "xls/dslx/frontend/pos.h"

namespace xls::dslx {

// Bindings (name references in the environment that map back to definition
// points in the AST) resolve to this BoundNode variant, which are all kinds of
// definitions.
using BoundNode =
    std::variant<EnumDef*, TypeAlias*, ConstantDef*, NameDef*, BuiltinNameDef*,
                 StructDef*, ProcDef*, Import*, UseTreeEntry*>;

// Returns a string, useful for reporting in error messages, for the type of the
// AST node contained inside of the given BoundNode variant; e.g. "Import".
std::string BoundNodeGetTypeString(const BoundNode& bn);

// Encodes ParseError data in a canonical way inside of an invalid argument
// error.
//
// When these propagate up to a Python boundary we throw them as exceptions
// using the data encoded in the absl::Status message. See ParseErrorGetData()
// for the utility used to extract the data from the error message text.
inline absl::Status ParseErrorStatus(const Span& span, std::string_view message,
                                     const FileTable& file_table) {
  return absl::InvalidArgumentError(
      absl::StrFormat("ParseError: %s %s", span.ToString(file_table), message));
}

inline absl::Status ParseNameErrorStatus(const Span& span,
                                         std::string_view name,
                                         const FileTable& file_table) {
  return ParseErrorStatus(
      span, absl::StrFormat("Cannot find a definition for name: \"%s\"", name),
      file_table);
}

// If `status` has a message containing a ParseNameErrorStatus payload as
// created above, extracts the name that the ParseNameError is referring to, or
// returns nullopt (e.g. if the status error code is not as expected, or it's
// not an appropriate error message).
std::optional<std::string_view> MaybeExtractParseNameError(
    const absl::Status& status);

struct PositionalErrorData {
  std::vector<Span> spans;
  std::string message;
  std::string error_type;

  std::string GetMessageWithType() const {
    return absl::StrCat(error_type, ": ", message);
  }

  bool operator==(const PositionalErrorData& other) const {
    return spans == other.spans && message == other.message &&
           error_type == other.error_type;
  }
};

// Returns parsed error data, or an error status if "status" is not of the
// special "positional error" format; e.g. of the formal generated by
// ParseErrorStatus() above.
absl::StatusOr<PositionalErrorData> GetPositionalErrorData(
    const absl::Status& status, std::optional<std::string_view> target_type,
    FileTable& file_table);

// Maps identifiers to the AST node that bound that identifier (also known as
// the lexical environment).
//
// The datatype is "stackable" so that we can easily take the bindings at a
// given point in the program (say in a function) and extend it with a new scope
// by stacking a fresh Bindings object on top (also sometimes referred to as a
// "scope chain"). For example:
//
//    Binding builtin_bindings;
//    builtin_bindings.Add("range", m.Make<BuiltinNameDef>("range"));
//
//    // Create a fresh scope, with no need to copy the builtin_bindings object.
//    Bindings function_bindings(&builtin_bindings);
//    XLS_ASSIGN_OR_RETURN(Function* f, ParseFunction(&function_bindings));
//
// We can do this because bindings are immutable and stack according to lexical
// scope; new bindings in the worst case only shadow previous bindings.
class Bindings {
 public:
  explicit Bindings(Bindings* parent = nullptr);

  // Too easy to confuse chaining a bindings object with its parent vs copy
  // constructing so we rely on an explicit Clone() call instead.
  Bindings(const Bindings& other) = delete;

  // Moving is ok though.
  Bindings(Bindings&& other) = default;
  Bindings& operator=(Bindings&& other) = default;

  void ConsumeLocalBindingsFrom(const Bindings& other) {
    for (const auto& [k, v] : other.local_bindings_) {
      local_bindings_[k] = v;
    }
  }

  // The "Cronus" method. This adds a child's bindings to this object, i.e., it
  // "commits" changes made in a child Bindings to this parent object.
  void ConsumeChild(const Bindings* child) {
    CHECK_EQ(child->parent_, this);
    ConsumeLocalBindingsFrom(*child);
  }

  // Returns whether there are any local bindings (i.e. bindings that are not
  // set in parent/ancestors).
  bool HasLocalBindings() const { return !local_bindings_.empty(); }

  // Adds a local binding.
  void Add(std::string name, BoundNode binding) {
    // When parsing the builtins module, the parser creates normal NameDefs
    // for what are really builtins. This is fine as long as we don't let these
    // NameDefs shadow the corresponding BuiltinNameDefs for the same entities.
    // Without this check, a builtin that uses another builtin in its signature
    // would cause problems in type concretization.
    if (IsBuiltinModule() && ResolveNode(name).has_value()) {
      return;
    }
    local_bindings_[std::move(name)] = binding;
  }

  void SetBuiltinModule(bool value) { builtin_module_ = value; }
  bool IsBuiltinModule() {
    return builtin_module_ ||
           (parent_ != nullptr && parent_->IsBuiltinModule());
  }

  // Returns whether the bindings are for a context that is inside a trait.
  bool IsInTrait() const {
    return in_trait_ || (parent_ != nullptr && parent_->IsInTrait());
  }

  // Returns whether `self` has been set to anything on this object or one of
  // its parents.
  bool HasSelf() const {
    return self_.has_value() || (parent_ != nullptr && parent_->HasSelf());
  }

  // Returns the `TypeAnnotation` for the struct that `self` indicates in the
  // current context. This only returns an annotation in an `impl`. In a trait
  // or other context, it returns `nullopt`.
  std::optional<TypeAnnotation*> GetImplSelf() const {
    if (!self_.has_value() && parent_ != nullptr) {
      return parent_->GetImplSelf();
    }
    if (self_.has_value() && std::holds_alternative<TypeAnnotation*>(*self_)) {
      return std::get<TypeAnnotation*>(*self_);
    }
    return std::nullopt;
  }

  // Indicates that `self` should be understood in this context as the object
  // implementing the trait whose name is `trait_name_def`. This means `self` is
  // valid but we don't have an actual `TypeAnnotation` for it.
  void AddTraitAsSelf(NameDef* trait_name_def) {
    CHECK(!HasSelf());
    self_ = trait_name_def;
    in_trait_ = true;
  }

  // Indicates that `self` should be understood in this context as some instance
  // of `struct_ref`.
  void AddStructAsSelf(TypeAnnotation* struct_ref) {
    CHECK(!HasSelf());
    self_ = struct_ref;
  }

  // fail! labels must be unique within a function.
  //
  // The labels are used in Verilog assertion labels, though they are given as
  // strings in the DSLX source.
  //
  // If a fail label is duplicated a parse error is returned.
  absl::Status AddFailLabel(const std::string& label, const Span& span,
                            const FileTable& file_table);

  // Returns the AST node bound to 'name'.
  std::optional<BoundNode> ResolveNode(std::string_view name) const {
    for (const Bindings* b = this; b != nullptr; b = b->parent_) {
      auto it = b->local_bindings_.find(name);
      if (it != b->local_bindings_.end()) {
        return it->second;
      }
    }
    return std::nullopt;
  }

  bool ResolveNodeIsTypeDefinition(std::string_view name) const {
    std::optional<BoundNode> bn = ResolveNode(name);
    if (!bn) {
      return false;
    }
    return std::holds_alternative<EnumDef*>(*bn) ||
           std::holds_alternative<TypeAlias*>(*bn) ||
           std::holds_alternative<StructDef*>(*bn) ||
           std::holds_alternative<ProcDef*>(*bn);
  }

  // As above, but flags a ParseError() if the binding cannot be resolved,
  // attributing the source of the binding resolution as span.
  absl::StatusOr<BoundNode> ResolveNodeOrError(
      std::string_view name, const Span& span,
      const FileTable& file_table) const {
    std::optional<BoundNode> result = ResolveNode(name);
    if (result.has_value()) {
      return *result;
    }
    return ParseNameErrorStatus(span, name, file_table);
  }

  // Resolves "name" as an AST binding and returns the associated NameDefNode.
  //
  // Returns nullopt if no AST node binding is found associated with "name".
  std::optional<AnyNameDef> ResolveNameOrNullopt(std::string_view name) const;

  // As above, but returns a ParseError status.
  absl::StatusOr<AnyNameDef> ResolveNameOrError(
      std::string_view name, const Span& span,
      const FileTable& file_table) const;

  // Returns whether there is an AST binding associated with "name".
  bool HasName(std::string_view name) const {
    return ResolveNode(name).has_value();
  }

  const absl::flat_hash_map<std::string, BoundNode>& local_bindings() const {
    return local_bindings_;
  }

  // Some properties, such as failure labels, are uniquified at a function
  // scope, so in parsing we mark which bindings are the "roll up point" for a
  // function scope so we can check for uniqueness there.
  void NoteFunctionScoped() {
    function_scoped_ = true;
    fail_labels_.emplace();
  }

  std::vector<std::string> GetLocalBindings() const {
    std::vector<std::string> result;
    for (const auto& [k, _] : local_bindings_) {
      result.push_back(k);
    }
    std::sort(result.begin(), result.end());
    return result;
  }

 private:
  Bindings* parent_;
  absl::flat_hash_map<std::string, BoundNode> local_bindings_;

  // Indicates that these bindings were created for a function scope -- this
  // helps us track properties that should be unique at function scope.
  bool function_scoped_ = false;

  // Note: only the function-scoped bindings will have fail_labels_.
  std::optional<absl::flat_hash_set<std::string>> fail_labels_;

  // This is a TypeAnnotation* when we are inside an impl; NameDef* when we are
  // inside a trait; nullopt otherwise.
  std::optional<std::variant<TypeAnnotation*, NameDef*>> self_;
  bool in_trait_ = false;
  bool builtin_module_ = false;
};

// Returns the name definition node (either builtin or user-defined) associated
// with the given binding data.
AnyNameDef BoundNodeToAnyNameDef(BoundNode bn);

// Returns the text span where the binding data is defined.
//
// For a builtin name definition, a "fake" span is returned (that spans no
// characters in the filename "<builtin>").
Span BoundNodeGetSpan(BoundNode bn, FileTable& file_table);

}  // namespace xls::dslx

#endif  // XLS_DSLX_FRONTEND_BINDINGS_H_
